import os
import sys
sys.path.append("..")
sys.path.append("../..")
from Dataloader.dataloader_utils import shuffle_data
from Dataloader.twitterloader import TwitterSet
import torch.nn as nn
from SentModel.Sent2Vec import W2VRDMVec
from PropModel.SeqPropagation import GRUModel
import torch
import math
from RumdetecFramework.AdverRumorFramework import EnhancedPropAdver

def obtain_general_set(tr_prefix, dev_prefix, te_prefix):
    tr_set = TwitterSet()
    tr_set.load_data_fast(data_prefix=tr_prefix, min_len=5)
    dev_set = TwitterSet()
    dev_set.load_data_fast(data_prefix=dev_prefix)
    te_set = TwitterSet()
    te_set.load_data_fast(data_prefix=te_prefix)
    return tr_set, dev_set, te_set

log_dir = str(__file__).split(".")[0]
if not os.path.exists(log_dir):
    os.system("mkdir %s"%log_dir)

for t in range(10):
    for i in range(5):
        tr, dev, te = obtain_general_set("../../data/twitter_tr%d"%i, "../../data/twitter_dev%d"%i, "../../data/twitter_te%d"%i)
        tr, dev, te = shuffle_data(tr, dev, te)
        lvec = W2VRDMVec("../../saved/glove_en/", 300, seg=None, emb_update=False)
        prop_share = GRUModel(300, 256, 1, 0.2)
        prop_specific = GRUModel(300, 256, 1, 0.2)
        cls = nn.Linear(256, 2)
        model = EnhancedPropAdver(lvec, prop_specific, prop_share, cls,
                         topic_label_num=5, prop_specific_only=True, batch_size=20, grad_accum_cnt=1)

        model.TopicPretrain(tr, dev, te)
        tr, dev, te = obtain_general_set("../../data/twitter_tr%d"%i, "../../data/twitter_dev%d"%i, "../../data/twitter_te%d"%i)
        test_event_name = te.data[te.data_ID[0]]['event']
        tr.filter_short_seq(min_len=5)
        tr.trim_long_seq(10)
        print("%s : (dev event)/(test event)/(train event) = %3d/%3d/%3d" % (
        test_event_name, len(dev), len(te), len(tr)))
        print("\n\n===========%s Train===========\n\n"%te.data[te.data_ID[0]]['event'])
        model.AdverTrain(tr, dev, te,
                        Unseen_every=10, lambda_1=0.9, lambda_2=0.9,
                        valid_every=100, max_epochs=10, lr_discount=1.0,
                        log_dir=log_dir, log_suffix=test_event_name, model_file="../../saved/%s_%s.pkl"%(log_dir, test_event_name))